import pandas as pd
import os
import statsmodels.api as sm
from datetime import datetime

# ---------------------- 配置参数（根据实际情况修改） ----------------------
BASE_DATA_PATH = "数据存储路径/"  # 基础数据根路径
FACTOR裂变_FOLDER = "因子裂变数据表/"  # 新表存储文件夹
today = datetime.now().strftime("%Y%m%d")  # 当日日期

# 需要处理的基础因子列表（根据你的实际因子名修改）
FACTORS = ['premium_rt','fixed_premium','curr_iss_amt',
           'bond_premium','excess_increase','bias_5','deviation_degree',
                'option_value','stock_wave','theory_premium','turnover_rt','ytm_rt']

# 中性化基准变量（通常为市值、行业等，根据你的数据列名修改）
# 若没有特定基准，可改为其他变量（如'price'）或删除相关逻辑
NEUTRALIZATION_BASE = 'curr_cap'  # 假设数据中有"market_cap"（市值）列curr_cap


# ---------------------- 核心函数：标准化因子（Z-score） ----------------------
def standardize_factors(df, factors):
    """
    对多个因子进行标准化处理（Z-score：(x - 均值)/标准差）
    新增列名格式：std_原因子名
    """
    df_standard = df.copy()
    for factor in factors:
        # 计算均值和标准差（排除空值）
        factor_mean = df_standard[factor].mean()
        factor_std = df_standard[factor].std()
        
        # 处理标准差为0的极端情况
        if factor_std == 0:
            df_standard[f'std_{factor}'] = 0.0
            print(f"⚠️ 因子[{factor}]所有值相同，标准化后均为0")
        else:
            df_standard[f'std_{factor}'] = (df_standard[factor] - factor_mean) / factor_std
        print(f"✅ 完成标准化：{factor} → std_{factor}")
    return df_standard


# ---------------------- 核心函数：中性化因子（回归残差法） ----------------------
def neutralize_factors(df, factors, base_col):
    """
    对多个因子进行中性化处理（去除基准变量影响）
    方法：以基准变量为自变量做线性回归，取残差作为中性化因子
    新增列名格式：neutral_原因子名
    """
    df_neutral = df.copy()
    # 确保基准变量无空值（中性化需要完整的基准数据）
    if df_neutral[base_col].isna().any():
        df_neutral = df_neutral.dropna(subset=[base_col])
        print(f"⚠️ 基准变量[{base_col}]存在空值，已剔除对应行（剩余{len(df_neutral)}条）")
    
    for factor in factors:
        # 准备回归数据（排除因子空值）
        regression_data = df_neutral.dropna(subset=[factor])
        X = regression_data[base_col]  # 自变量：基准变量
        y = regression_data[factor]    # 因变量：待中性化因子
        X = sm.add_constant(X)  # 添加常数项
        
        # 线性回归并取残差（残差 = 实际值 - 预测值，即去除基准影响后的因子）
        model = sm.OLS(y, X).fit()
        df_neutral.loc[regression_data.index, f'neutral_{factor}'] = model.resid
        print(f"✅ 完成中性化：{factor} → neutral_{factor}")
    return df_neutral


# ---------------------- 主流程：读取数据→处理因子→保存结果 ----------------------
def process_and_save_factors():
    # 1. 构建目标文件夹（自动创建）
    target_folder = os.path.join(BASE_DATA_PATH, FACTOR裂变_FOLDER)
    os.makedirs(target_folder, exist_ok=True)
    print(f"📂 新表将保存至：{target_folder}\n")

    # 2. 读取原始数据
    original_file = f"{BASE_DATA_PATH}/{today}_全量数据.xlsx"
    if not os.path.exists(original_file):
        raise FileNotFoundError(f"❌ 未找到原始数据文件：{original_file}")
    df = pd.read_excel(original_file)
    print(f"📊 已读取原始数据（共{len(df)}条记录）")

    # 3. 检查基础因子和基准变量是否存在
    missing_factors = [f for f in FACTORS if f not in df.columns]
    if missing_factors:
        raise ValueError(f"❌ 原始数据中缺少因子：{missing_factors}")
    if NEUTRALIZATION_BASE not in df.columns:
        raise ValueError(f"❌ 原始数据中缺少中性化基准变量：{NEUTRALIZATION_BASE}")

    # 4. 初步清洗（去除所有因子的空值）
    df_clean = df.dropna(subset=FACTORS).copy()
    print(f"🧹 清洗后数据量：{len(df_clean)}条（剔除因子空值{len(df)-len(df_clean)}条）\n")

    # 5. 批量处理因子（先标准化，再中性化）
    df_processed = standardize_factors(df_clean, FACTORS)  # 标准化
    df_processed = neutralize_factors(df_processed, FACTORS, NEUTRALIZATION_BASE)  # 中性化

    # 6. 保存结果（包含原始因子+标准化因子+中性化因子）
    output_file = os.path.join(target_folder, f"{today}_裂变因子全量数据.xlsx")
    df_processed.to_excel(output_file, index=False)
    print(f"\n💾 处理完成！结果已保存至：{output_file}")

    return df_processed


# ---------------------- 执行主流程 ----------------------
if __name__ == "__main__":
    result_df = process_and_save_factors()
    
    # 预览部分结果（原始因子+标准化+中性化各1个示例）
    if FACTORS:
        sample_factor = FACTORS[0]
        print("\n📌 结果预览（前5行）：")
        print(result_df[[sample_factor, f'std_{sample_factor}', f'neutral_{sample_factor}']].head())